"""
数据存储模块
使用SQLite存储进程内存监控数据
"""
import sqlite3
import os
from datetime import datetime
from typing import Optional, List, Tuple


class Database:
    """数据库操作类"""
    
    def __init__(self, db_path: str = "memory_monitor.db"):
        """
        初始化数据库连接
        
        Args:
            db_path: 数据库文件路径
        """
        self.db_path = db_path
        self._init_database()
    
    def _init_database(self):
        """初始化数据库表结构"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建监控记录表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS memory_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                process_name TEXT,
                process_id INTEGER,
                memory_mb REAL NOT NULL,
                memory_percent REAL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # 创建索引以提高查询性能
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_timestamp 
            ON memory_records(timestamp)
        """)
        
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_process 
            ON memory_records(process_id, timestamp)
        """)
        
        conn.commit()
        conn.close()
    
    def insert_record(self, process_name: Optional[str], process_id: int, 
                     memory_mb: float, memory_percent: Optional[float] = None):
        """
        插入一条监控记录
        
        Args:
            process_name: 进程名称
            process_id: 进程ID
            memory_mb: 内存占用（MB）
            memory_percent: 内存占用百分比
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        timestamp = datetime.now().isoformat()
        
        cursor.execute("""
            INSERT INTO memory_records 
            (timestamp, process_name, process_id, memory_mb, memory_percent)
            VALUES (?, ?, ?, ?, ?)
        """, (timestamp, process_name, process_id, memory_mb, memory_percent))
        
        conn.commit()
        conn.close()
    
    def get_records(self, process_id: Optional[int] = None, 
                   start_time: Optional[str] = None,
                   end_time: Optional[str] = None) -> List[Tuple]:
        """
        查询监控记录
        
        Args:
            process_id: 进程ID，如果为None则查询所有进程
            start_time: 开始时间（ISO格式字符串）
            end_time: 结束时间（ISO格式字符串）
        
        Returns:
            记录列表，每条记录包含 (timestamp, process_name, process_id, memory_mb, memory_percent)
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT timestamp, process_name, process_id, memory_mb, memory_percent FROM memory_records WHERE 1=1"
        params = []
        
        if process_id is not None:
            query += " AND process_id = ?"
            params.append(process_id)
        
        if start_time:
            query += " AND timestamp >= ?"
            params.append(start_time)
        
        if end_time:
            query += " AND timestamp <= ?"
            params.append(end_time)
        
        query += " ORDER BY timestamp ASC"
        
        cursor.execute(query, params)
        records = cursor.fetchall()
        conn.close()
        
        return records
    
    def clear_records(self, process_id: Optional[int] = None):
        """
        清空记录
        
        Args:
            process_id: 进程ID，如果为None则清空所有记录
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if process_id is not None:
            cursor.execute("DELETE FROM memory_records WHERE process_id = ?", (process_id,))
        else:
            cursor.execute("DELETE FROM memory_records")
        
        conn.commit()
        conn.close()
    
    def get_statistics(self, process_id: Optional[int] = None) -> dict:
        """
        获取统计信息
        
        Args:
            process_id: 进程ID，如果为None则统计所有进程
        
        Returns:
            包含最大值、最小值、平均值等统计信息的字典
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT MAX(memory_mb), MIN(memory_mb), AVG(memory_mb), COUNT(*) FROM memory_records WHERE 1=1"
        params = []
        
        if process_id is not None:
            query += " AND process_id = ?"
            params.append(process_id)
        
        cursor.execute(query, params)
        result = cursor.fetchone()
        conn.close()
        
        return {
            'max_memory_mb': result[0] if result[0] else 0,
            'min_memory_mb': result[1] if result[1] else 0,
            'avg_memory_mb': result[2] if result[2] else 0,
            'record_count': result[3] if result[3] else 0
        }

